#include <stdio.h>
#include <stdlib.h>
#include "image.h"
#include "minst_data_set.h"

static minst_elem_t minst_elem_ = {0};

int create_minst_data_set(char *data_set, char *data_label)
{
    int magic;
    int file_size;
    minst_elem_t minst_elem = {0};
    char *image_ptr;
    char byte_buf[4];
    char tmp;
    int i;
    FILE *data_set_fp;
    FILE *data_label_fp;
    data_label_fp = fopen(data_label, "r");
    // fseek(data_label_fp, 8, SEEK_SET);
    fread(byte_buf, sizeof(int), 1, data_label_fp);
    fread(byte_buf, sizeof(int), 1, data_label_fp);
    data_set_fp = fopen(data_set, "r");
    fseek(data_set_fp, 0, SEEK_END);
    file_size = ftell(data_set_fp);
    rewind(data_set_fp);
    fread(byte_buf, sizeof(int), 1, data_set_fp);
    for (i = 0; i < 2; i++) {
        tmp = byte_buf[i];
        byte_buf[i] = byte_buf[3 - i];
        byte_buf[3 - i] = tmp;
    }
    magic = *((int*)byte_buf);
    file_size -= 4;
    fread(byte_buf, sizeof(int), 1, data_set_fp);
    for (i = 0; i < 2; i++) {
        tmp = byte_buf[i];
        byte_buf[i] = byte_buf[3 - i];
        byte_buf[3 - i] = tmp;
    }
    minst_elem.count = (int)(*((int*)byte_buf));
    file_size -= 4;
    fread(byte_buf, sizeof(int), 1, data_set_fp);
    for (i = 0; i < 2; i++) {
        tmp = byte_buf[i];
        byte_buf[i] = byte_buf[3 - i];
        byte_buf[3 - i] = tmp;
    }
    minst_elem.width = (int)(*((int*)byte_buf));
    file_size -= 4;
    fread(byte_buf, sizeof(int), 1, data_set_fp);
    for (i = 0; i < 2; i++) {
        tmp = byte_buf[i];
        byte_buf[i] = byte_buf[3 - i];
        byte_buf[3 - i] = tmp;
    }
    minst_elem.height = (int)(*((int*)byte_buf));
    file_size -= 4;

    minst_elem.image = calloc(minst_elem.count * minst_elem.width
        *minst_elem.height, sizeof(char));
    image_ptr = minst_elem.image;
    while (file_size)
    {
        fread(byte_buf, sizeof(int), 1, data_set_fp);
        // image_ptr[3] = byte_buf[0];
        // image_ptr[2] = byte_buf[1];
        // image_ptr[1] = byte_buf[2];
        // image_ptr[0] = byte_buf[3];
        image_ptr[0] = byte_buf[0];
        image_ptr[1] = byte_buf[1];
        image_ptr[2] = byte_buf[2];
        image_ptr[3] = byte_buf[3];
        file_size -= 4;
        image_ptr += 4;
    }
    fclose(data_set_fp);
    int index;
    minst_elem.label = calloc(minst_elem.count, sizeof(char));
    image_ptr = minst_elem.label;
    for (index = 0; index < minst_elem.count; index += 4) {
        fread(byte_buf, sizeof(int), 1, data_label_fp);
        image_ptr[0] = byte_buf[0];
        image_ptr[1] = byte_buf[1];
        image_ptr[2] = byte_buf[2];
        image_ptr[3] = byte_buf[3];
        // image_ptr[3] = byte_buf[0];
        // image_ptr[2] = byte_buf[1];
        // image_ptr[1] = byte_buf[2];
        // image_ptr[0] = byte_buf[3];
        image_ptr += 4;
    }
    fclose(data_label_fp);
    index = 0;
    printf("magic %x count %d width %d height %d %d lable %d\n", 
        magic, minst_elem.count, minst_elem.width,
         minst_elem.height, index, 
         minst_elem.label[index]);
    // data_set_fp = fopen("test.raw", "w");
    // fwrite(minst_elem.image + index * minst_elem.width * minst_elem.height, 
    //     sizeof(char), minst_elem.width * minst_elem.height, data_set_fp);
    // tmp = 0;
    // for (i = 0; i < minst_elem.width * minst_elem.height / 2; i++) {
    //     fwrite(&tmp, sizeof(char), 1, data_set_fp);
    // }
    // fclose(data_set_fp);

    minst_elem_ = minst_elem;
    return 0;
}

data batch_load_minst_data(int batch, int start, int rand)
{
    data d = {0};
    matrix y;
    matrix X;
    int i,j,k;

    if (start < minst_elem_.count) {
        printf("batch load index %d label %d count %d\n", 
            start, minst_elem_.label[start],minst_elem_.count);
    } else
    {
        start = 0;
        printf("batch load index %d label %d count %d\n", 
            start, minst_elem_.label[start],minst_elem_.count);
    }

    y = make_matrix(batch, 10);
    X = make_matrix(batch, minst_elem_.width * minst_elem_.height * 1);

    for (k = 0; k < batch; k++) {
        for (i = 0; i < minst_elem_.height; i++) {
            for (j = 0; j < minst_elem_.width; j++) {
                // printf("%u ", minst_elem_.image[i * minst_elem_.width + j]);
                // im.data[(i + 2) * 32 + j + 2] = (float)minst_elem_.image[i * minst_elem_.width + j] / 255.0;
                X.vals[k][i * minst_elem_.width + j] = (float)minst_elem_.image[(start + k) * minst_elem_.width * minst_elem_.height
                     + i * minst_elem_.width + j] / 255.0;
                // printf("%f ", X.vals[k][i * minst_elem_.width + j]);
            }
            // printf("\n");
        }
        y.vals[k][minst_elem_.label[start + k]] = 1;
    }
    printf("\n");
    // print_matrix(y);
    d.shallow = 0;
    d.X = X;
    d.y = y;

    return d;
}

void destroy_minst_data_set()
{
    if (minst_elem_.label) {
        free(minst_elem_.label);
    }    
    if (minst_elem_.image) {
        free(minst_elem_.image);
    }

}